package interval

import (
	"math"

	"golang.org/x/exp/constraints"
)

type Fenwick[T constraints.Integer] struct {
	size int
	a    []T
}

func NewFenwick[T constraints.Integer](n int) *Fenwick[T] {
	return &Fenwick[T]{
		size: n,
		a:    make([]T, n),
	}
}

func (fwk *Fenwick[T]) Add(x int, v T) {
	for i := x + 1; i <= fwk.size; i += i & -i {
		fwk.a[i-1] = fwk.a[i-1] + v
	}
}

func (fwk *Fenwick[T]) Sum(x int) T {
	var ans T
	for i := x; i > 0; i -= i & -i {
		ans = ans + fwk.a[i-1]
	}
	return ans
}

func (fwk *Fenwick[T]) RangeSum(l, r int) T {
	return fwk.Sum(r) - fwk.Sum(l)
}

func (fwk *Fenwick[T]) Select(k T) int {
	x := 0
	var cur T
	for i := 1 << int(math.Log2(float64(fwk.size))); i != 0; i /= 2 {
		if x+i <= fwk.size && cur+fwk.a[x+i-1] <= k {
			x += i
			cur += fwk.a[x-1]
		}
	}
	return x
}
